import os
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import torch
import cv2
import argparse
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import auxil
import models.presnet as PYRM
from main import predict

def production_data_valtespre(c, tes_num, half_s, image_3d_mat, flag='Tes'):
    ##Testing set(spatial)
    tes_num = np.array(tes_num)
    idx_2d_tes = np.zeros([tes_num.shape[0], 2]).astype(int)
    idx_2d_tes[:, 0] = np.floor(tes_num / c)
    idx_2d_tes[:, 1] = tes_num + 1 - c * idx_2d_tes[:, 0] - 1
    # neibour area(2*half_s+1)
    patch_size = 2 * half_s + 1
    tes_spat = np.zeros([tes_num.shape[0], patch_size, patch_size, image_3d_mat.shape[2]])
    for i in range(idx_2d_tes.shape[0]):
        # 图像扩展
        row = idx_2d_tes[i, 0] + half_s
        col = idx_2d_tes[i, 1] + half_s
        tes_spat[i, :, :, :] = image_3d_mat[(row - half_s):row + half_s + 1,
                               (col - half_s):col + half_s + 1, :]

    print('tes_spat:', tes_spat.shape)
    print('{} Spatial dataset preparation Finished!'.format(flag))
    return tes_spat, tes_num

## 指定显卡
# USE_GPU=True
# if USE_GPU:
#     os.environ["CUDA_VISIBLE_DEVICES"] = "3"
# else:
#     device=torch.device('cpu')


parser = argparse.ArgumentParser(description='PyTorch DCNNs Training')
parser.add_argument('--components', default=None, type=int, help='dimensionality reduction')
parser.add_argument('--dataset', default='IP', type=str, help='dataset (options: IP, UP, SV, KSC)')
parser.add_argument('--te_bsize', default=1000, type=int, help='mini-batch test size (default: 1000)')
parser.add_argument('--spatialsize', dest='spatialsize', default=7, type=int, help='spatial-spectral patch dimension')
parser.add_argument('--inplanes', dest='inplanes', default=16, type=int, help='bands before blocks')
parser.add_argument('--no-bottleneck', dest='bottleneck', action='store_false',
                    help='to use basicblock (default: bottleneck)')
parser.add_argument('--depth', default=32, type=int, help='depth of the network (default: 32)')
parser.add_argument('--alpha', default=50, type=int, help='number of new channel increases per depth (default: 12)')

parser.set_defaults(bottleneck=True)
args = parser.parse_args()

if args.spatialsize < 9:
    avgpoosize = 1
elif args.spatialsize <= 11:
    avgpoosize = 2
elif args.spatialsize == 15:
    avgpoosize = 3
elif args.spatialsize == 19:
    avgpoosize = 4
elif args.spatialsize == 21:
    avgpoosize = 5
elif args.spatialsize == 27:
    avgpoosize = 6
elif args.spatialsize == 29:
    avgpoosize = 7
else:
    print("spatialsize no tested")

###################### 加载数据, 模型及设置参数及全局变量 ####################

#每块样本数
BATCH=50000

data, labels, num_classes = auxil.loadData(args.dataset, num_components=args.components)

margin = int((args.spatialsize-1)/2)

zeroPaddedX = auxil.padWithZeros(data, margin=margin)##零值填充

label_mat=labels.reshape(-1,1)

r,c,n_bands=data.shape[0],data.shape[1],data.shape[-1]

del data

pre_num=np.array(list(range(0,r*c)))

###### 网络相关情况设置

model = PYRM.pResNet(args.depth, args.alpha, num_classes, n_bands, avgpoosize, args.inplanes, bottleneck=args.bottleneck) # for PyramidNet
model = torch.nn.DataParallel(model)

use_cuda = torch.cuda.is_available()
if use_cuda: model = model.cuda()

checkpoint = torch.load("best_model_"+str(args.dataset)+".pth.tar")
model.load_state_dict(checkpoint['state_dict'])

#net=net.module#if use DataParallel

criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失函数

###### 标签相关情况设置
y_disp=np.zeros([r*c])

y_disp_all=y_disp.copy()

start=0
end=np.min([start+BATCH,pre_num.shape[0]])

###################################### 预测 #######################################

part_num=int(pre_num.shape[0]/BATCH)+1

print('需要分成{}块来预测'.format(part_num))

for i in range(0,part_num):

    ###################### label ######################

    pre_num_part=pre_num[start:end]

    y_pre=label_mat[pre_num_part,0]#include background

    pre_YY = torch.from_numpy(np.ones([y_pre.shape[0]]))

    ######################  data ######################

    pre_spat, pre_num_part = production_data_valtespre(c,pre_num_part, margin, zeroPaddedX, flag='Pre')

    pre_XX_spat=torch.from_numpy(pre_spat.transpose(0, 3, 1, 2))

    del pre_spat

    ######### 推断，预测集 #########

    pre_XX_spat = pre_XX_spat.type(torch.FloatTensor)

    pre_dataset=TensorDataset(pre_XX_spat,pre_YY)
    pre_loader=DataLoader(pre_dataset,batch_size=args.te_bsize)

    y_pred_pre=np.argmax(predict(pre_loader, model, criterion, use_cuda),axis=1)+1

    print('第{}块数据预测完成！！！'.format(i))

    y_disp_all[pre_num_part]=y_pred_pre

    start=end
    end=np.min([start+BATCH,pre_num.shape[0]])

######################  展示   ##########################

#np.save('y_pre_before10w.npy',y_disp_all)

# plt.xlabel('All image')
# plt.imshow(y_disp_all.reshape(r,c),cmap='jet')
# plt.xticks([])
# plt.yticks([])
#
# plt.show()

#####################  保存预测全图 #####################

#保存生成的预测图便于后期上色

#cv2.imwrite('DPRN_all_'+str(args.dataset)+'.png', y_disp_all.reshape(r,c))

plt.subplots(figsize=[10,10])
a1=plt.imshow(y_disp_all.reshape(r,c),cmap='jet')
plt.xticks([])
plt.yticks([])
plt.savefig('DPRN_all_'+str(args.dataset)+'.png',dpi=600,bbox_inches='tight')
#plt.show()

print('结果图展示 & 保存阶段完成！！')